# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import ujson as json
import functools
import datetime

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.binance.book_builder import BinanceBookBuilder
from coin.exchange.binance.kr_rest.product import BinanceProduct


class BinanceHandlerBase(WebsocketHandler):
  BookBuilderType = None
  TraitsType = None

  def on_rest_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())
    symbol = data['symbol']
    if not self.is_subscribed('book', symbol):
      return
    last_update_id = data['data']['lastUpdateId']

    if symbol not in self._books:
      book = self.BookBuilderType(symbol,
                                  last_update_id,
                                  book_builder_observer=self._book_builder_observer)
      self._books[symbol] = book
    else:
      book = self._books[symbol]
      book.reset_last_update_id(last_update_id)
      book.clear_book()

    bids = data['data']['bids']
    asks = data['data']['asks']
    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self.ready = True
    self.publish_event('book', book.symbol, book)
    product = self.TraitsType.get_product_from_symbol(symbol, record.timestamp)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)

  def on_ws_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())
    if 'data' in data:
      data = data['data']
    event = data['e']
    symbol = data['s']

    if symbol in self._books:
      book = self._books[symbol]
      assert event == 'depthUpdate'
      first_update_id = data['U']
      last_update_id = data['u']
      bids = data['b']
      asks = data['a']
      book.update(first_update_id, last_update_id, bids, asks)
      book.update_timestamp(record.timestamp)
      book.update_topic_id(record.topic_id)
      product = self.TraitsType.get_product_from_symbol(symbol, record.timestamp)
      if self._feed_checker:
        self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
      self.publish_event('book', book.symbol, book)

  def on_ws_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())
    if 'data' in data:
      data = data['data']
    event = data['e']
    assert event in ['trade', 'aggTrade']
    symbol = data['s']
    if not self.is_subscribed('trade', symbol):
      return
    try:
      tsdiff_sec = (data['E'] - data['T']) * 1e-3
      if tsdiff_sec > 60:
        print("binance T ignored due to E time - T time diff : "
              f'{tsdiff_sec} = {data["E"] * 1e-3} - {data["T"] * 1e-3}')
        return
    except Exception:
      pass
    price = float(data['p'])
    qty = float(data['q'])
    side = TradeSide.TRADE_SELL_SIDE if data['m'] else TradeSide.TRADE_BUY_SIDE
    trade_event = create_market_trade_event(record.timestamp,
                                            symbol,
                                            price,
                                            qty,
                                            side,
                                            topic_id=record.topic_id)
    product = self.TraitsType.get_product_from_symbol(symbol, record.timestamp)
    if self._feed_checker:
      self._feed_checker.on_trade(product=product,
                                  trade_event=trade_event,
                                  timestamp=record.timestamp)
    self.publish_event('trade', symbol, trade_event)

  def on_l1_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict) and not isinstance(data, list):
      data = json.loads(data.tobytes())

    if 'data' in data:
      data = data['data']
    symbol = data['s']
    if not self.is_subscribed('book', symbol):
      return

    if symbol not in self._books:
      last_update_id = data['u']
      book = self.BookBuilderType(symbol,
                                  last_update_id,
                                  book_builder_observer=self._book_builder_observer)
      self._books[symbol] = book
    else:
      book = self._books[symbol]
      book.clear_book()

    bids = [[data['b'], data['B']]]
    asks = [[data['a'], data['A']]]

    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)

    self.ready = True
    self.publish_event('book', book.symbol, book)

    product = self.TraitsType.get_product_from_symbol(symbol, record.timestamp)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)


NANOSEC_ONE_DAY = 24 * 60 * 60 * 1e9


class SpotTraits(object):
  @classmethod
  def get_product_from_symbol(cls, symbol, timestamp):
    as_of_timestamp = timestamp // NANOSEC_ONE_DAY * NANOSEC_ONE_DAY
    return SpotTraits._get_product_from_symbol(symbol, as_of_timestamp)

  @classmethod
  @functools.lru_cache(maxsize=128)
  def _get_product_from_symbol(cls, symbol, as_of_timestamp):
    current_datetime = datetime.datetime.fromtimestamp(as_of_timestamp / 1e9)
    return BinanceProduct.FromStrNativeProduct(symbol, current_datetime)


class BinanceHandler(BinanceHandlerBase):
  BookBuilderType = BinanceBookBuilder
  TraitsType = SpotTraits
